﻿#pragma once
#include <memory>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
#include <iostream>
// #include <omp.h>

#include <opencv2/opencv.hpp>
#include <yaml-cpp/yaml.h>

#include "visualize.h"
#include "model_struct.h"

class LogInfo;

namespace ai {
class ModelConfig;

/*
 * @brief
 * Abstraction of preprocessing operation class
 * */
class Transform {
 public:
  Transform() {}
	virtual ~Transform() {spdlog::get("logger")->info("Transform::~Transform");};
  virtual void init(const YAML::Node& item) = 0;
  /*
   * @brief
   * This method executes preprocessing operation on image matrix,
   * result will be returned at second parameter.
   * @param im: single image matrix to be preprocessed
   * @param data: the raw data of single image matrix after preprocessed
   * @return true if transform successfully
   * */
  virtual bool run(cv::Mat& im, stream::ImageBlob& data, ModelConfig* cfg) = 0;
};

/*
 * @brief
 * This class execute normalization operation on image matrix
 * */
class Normalize : public Transform {
 public:
  Normalize() {}
	virtual ~Normalize() {spdlog::get("logger")->info("Normalize::~Normalize");};
  virtual void init(const YAML::Node& item) {
    if (item["scale"].IsDefined()) { scale = item["scale"].as<int>();} 
    // 默认为 -1 表示在Normalize中只做赋值操作
    else {scale = -1;}
    if (item["mean"].IsDefined()) { mean_ = item["mean"].as<std::vector<float>>();} 
    // else {mean_ = {0,0,0};}
    if (item["std"].IsDefined()) { std_ = item["std"].as<std::vector<float>>();} 
    // else {std_ = {1,1,1};}

    if (item["min_val"].IsDefined()) {
      min_val_ = item["min_val"].as<std::vector<float>>();
    } else {
      // min_val_ = std::vector<float>(mean_.size(), 0.);
    }
    if (item["max_val"].IsDefined()) {
      max_val_ = item["max_val"].as<std::vector<float>>();
    } else {
      // max_val_ = std::vector<float>(mean_.size(), 255.);
    }
  }

  virtual bool run(cv::Mat& im, stream::ImageBlob& data, ModelConfig* cfg);

 private:
  int scale;
  std::vector<float> mean_;
  std::vector<float> std_;
  std::vector<float> min_val_;
  std::vector<float> max_val_;
};

/*
 * @brief
 * This class execute resize by short operation on image matrix. At first, it resizes
 * the short side of image matrix to specified length. Accordingly, the long side
 * will be resized in the same proportion. If new length of long side exceeds max
 * size, the long size will be resized to max size, and the short size will be
 * resized in the same proportion
 * */
class ResizeByShort : public Transform {
 public:
    ResizeByShort() {}
    virtual ~ResizeByShort() {spdlog::get("logger")->info("ResizeByShort::~ResizeByShort");};  
    virtual void init(const YAML::Node& item) {
    short_size_ = item["short_size"].as<int>();
    if (item["max_size"].IsDefined()) {
      max_size_ = item["max_size"].as<int>();
    } else {
      max_size_ = -1;
    }
  }
  virtual bool run(cv::Mat& im, stream::ImageBlob& data, ModelConfig* cfg);

 private:
  float GenerateScale(const cv::Mat& im);
  int short_size_;
  int max_size_;
};

/*
 * @brief
 * This class execute resize by long operation on image matrix. At first, it resizes
 * the long side of image matrix to specified length. Accordingly, the short side
 * will be resized in the same proportion.
 * */
class ResizeByLong : public Transform {
 public:
  ResizeByLong() {}
  virtual ~ResizeByLong() {spdlog::get("logger")->info("ResizeByLong::~ResizeByLong");};  
  virtual void init(const YAML::Node& item) {
    if (item["max_size"].IsDefined()) {
      max_size = item["max_size"].as<int>();
    }
    if (item["fix_size"].IsDefined()) {
      fix_size = item["fix_size"].as<int>();
    }
    if ((max_size > 0 && fix_size > 0) || (max_size == 0 && fix_size == 0)) { 
      std::cerr << "param ResizeByLong error." << std::endl;
      exit(0);
    }
  }
  virtual bool run(cv::Mat& im, stream::ImageBlob& data, ModelConfig* cfg);

 private:
  int max_size = 0;   // 最大边长 如DB 超过最大边长才需要resize
  int fix_size = 0;   // 固定边长 如目标检测 尺寸固定
};

/*
 * @brief
 * This class execute resize operation on image matrix. It resizes width and height
 * to specified length.
 * */
class Resize : public Transform {
 public:
  Resize() {}
  virtual ~Resize() {spdlog::get("logger")->info("Resize::~Resize");};  
  virtual void init(const YAML::Node& item) {
    interp_ = "LINEAR";
    if (item["interp"].IsDefined()) {
      interp_ = item["interp"].as<std::string>();
    }
    if (item["target_size"].IsDefined()) {
      if (item["target_size"].IsScalar()) {
        target_size.push_back(item["target_size"].as<int>());
        target_size.push_back(item["target_size"].as<int>());
      } else if (item["target_size"].IsSequence()) {
        target_size = item["target_size"].as<std::vector<int>>();
      }
    }
    else if (item["fix_size"].IsDefined()) {
      if (item["fix_size"].IsScalar()) {
        fix_size.push_back(item["fix_size"].as<int>());
        fix_size.push_back(item["fix_size"].as<int>());
      }
      else if (item["fix_size"].IsSequence()) {
        fix_size = item["fix_size"].as<std::vector<int>>();
      }
    }
    if (target_size.size() < 2 && fix_size.size() < 2) {
      std::cerr << "[Resize] target_size should greater than 0" << std::endl;
      exit(-1);
    }
  }
  virtual bool run(cv::Mat& im, stream::ImageBlob& data, ModelConfig* cfg);

 private:
  std::vector<int> target_size;
  std::vector<int> fix_size;
  std::string interp_;
};

class ResizeText : public Transform {
 public:
  ResizeText() {}
  virtual ~ResizeText() {spdlog::get("logger")->info("ResizeText::~ResizeText");}; 
   virtual void init(const YAML::Node& item) {
    interp_ = "LINEAR";
    if (item["interp"].IsDefined()) {
      interp_ = item["interp"].as<std::string>();
    }
    std_height = 32;
    if (item["fix_height"].IsDefined()) {
      std_height = item["fix_height"].as<int>();
    }
  }
  virtual bool run(cv::Mat& im, stream::ImageBlob& data, ModelConfig* cfg);
 
 private:
  std::string interp_;
  int std_height;
};


/*
 * @brief
 * This class execute center crop operation on image matrix. It crops the center
 * of image matrix accroding to specified size.
 * */
class CenterCrop : public Transform {
 public:
  CenterCrop() {}
  virtual ~CenterCrop() {spdlog::get("logger")->info("CenterCrop::~CenterCrop");}; 
  virtual void init(const YAML::Node& item) {
    if (item["crop_size"].IsScalar()) {
      height_ = item["crop_size"].as<int>();
      width_ = item["crop_size"].as<int>();
    } else if (item["crop_size"].IsSequence()) {
      std::vector<int> crop_size = item["crop_size"].as<std::vector<int>>();
      width_ = crop_size[0];
      height_ = crop_size[1];
    }
  }
  virtual bool run(cv::Mat& im, stream::ImageBlob& data, ModelConfig* cfg);

 private:
  int height_;
  int width_;
};

/*
 * @brief
 * This class execute padding operation on image matrix. It makes border on edge
 * of image matrix.
 * */
class Padding : public Transform {
 public:
  Padding() {}
  virtual ~Padding() {spdlog::get("logger")->info("Padding::~Padding");}; 
 
  virtual void init(const YAML::Node& item) {
    if (item["coarsest_stride"].IsDefined()) {
      coarsest_stride_ = item["coarsest_stride"].as<int>();
      if (coarsest_stride_ < 1) {
        std::cerr << "[Padding] coarest_stride should greater than 0"
                  << std::endl;
        exit(-1);
      }
    }
    if (item["target_size"].IsDefined()) {
      if (item["target_size"].IsScalar()) {
        width_ = item["target_size"].as<int>();
        height_ = item["target_size"].as<int>();
      } else if (item["target_size"].IsSequence()) {
        width_ = item["target_size"].as<std::vector<int>>()[0];
        height_ = item["target_size"].as<std::vector<int>>()[1];
      }
    }
    if (item["im_padding_value"].IsDefined()) {
      im_value_ = item["im_padding_value"].as<std::vector<float>>();
    } else {
      im_value_ = {0, 0, 0};
    }
  }
  virtual bool run(cv::Mat& im, stream::ImageBlob& data, ModelConfig* cfg);
  virtual void GeneralPadding(cv::Mat& im,
                              const std::vector<float> &padding_val,
                              int padding_w, int padding_h);
  virtual void MultichannelPadding(cv::Mat& im,
                                   const std::vector<float> &padding_val,
                                   int padding_w, int padding_h);

 private:
  int coarsest_stride_ = -1;
  int width_ = 0;
  int height_ = 0;
  std::vector<float> im_value_;
};

/*
 * @brief
 * This class execute clip operation on image matrix
 * */
class Clip : public Transform {
 public:
  Clip() {}
  virtual ~Clip() {spdlog::get("logger")->info("Clip::~Clip");}; 
 
  virtual void init(const YAML::Node& item) {
    min_val_ = item["min_val"].as<std::vector<float>>();
    max_val_ = item["max_val"].as<std::vector<float>>();
  }

  virtual bool run(cv::Mat& im, stream::ImageBlob& data, ModelConfig* cfg);

 private:
  std::vector<float> min_val_;
  std::vector<float> max_val_;
};

/*
 * @brief
 * This class is transform operations manager. It stores all neccessary
 * transform operations and run them in correct order.
 * */

class ModelConfig;
class Transforms {
public:
  Transforms(LogInfo* li);
  ~Transforms();
  int init(const YAML::Node& node);
  Transform* CreateTransform(const std::string& name);
  bool run(cv::Mat& im, stream::ImageBlob& data, ModelConfig* cfg);

private:
  LogInfo* log_ifo;
  std::vector<Transform*> transforms;
};

}  // namespace PaddleX
